'''
Created on Jan 13, 2011

A place to put the Markov chain builder stuff.

@author: matt
'''

import sys, re, random
from heavenlyfodder.util.text import TextAtom
#import heavenlyfodder.util.text.TextAtom as TextAtom
#import heavenlyfodder.util.text.TextGroup as TextGroup
from heavenlyfodder.util.text import TextGroup

# TODO: We don't seem to be using this anymore; remove it if possible
#def to_text_atoms (input_text, delimiter, source_name):
#    """
#    Splits the input text into atoms, using the specified delimiter to
#    determine where one atom ends and the next one begins
#    """
#    
#    text_atoms = []
#    
#    textual_units = input_text.split(delimiter)
#    for unit in textual_units:
#        # Note that we're using the default (empty) delimiter
#        text_atoms.append(TextAtom(unit, source_name))
#        
#    return text_atoms

def atom_list_to_atom(atom_list, delimiter=' '):
    """
    Concatenates all the words in atom_list into a single string, with
    delimiter between each atom
    """
    
    if not atom_list:
        return None
    
    result = ''
    list_pos = 0
    for atom in atom_list:
        result += atom.original_text
        # This ensures that we don't put a delimiter on the end of the list
        # (we only want it between elements)
        if list_pos < len(atom_list) - 1:
            result += delimiter
        list_pos += 1 
            
#    return result
    # All the constituent atoms should be from the same source, so we use the
    # source field of the first element
    return TextAtom(result, atom_list[0].source, delimiter)
    

# TODO: modify this to allow the caller to specify which characters are to be
# used for ending a run, separating words etc. - OR - allow the caller to
# specify a class that provides methods that TextualUnitSlicer can call to 
# determine these things.
class TextualUnitSlicer(object):
    """
    Uses the iterator methods to return lists of words from input_text in a
    format suitable for use in building a markov chain.  This class currently
    specializes in chopping up short, self-contained bits of text (e.g.,
    film titles).
    """
    
    # TODO: Pick up the TextAtom migration HERE (then go back to the bottom of
    # this file, where TextualUnitSlicer is used)
    
    #TODO: when upgrading to Textatoms: instead of input_text being a big
    # string, have it be a TextAtom - the next() method might also *return*
    # Textatoms (instead of strings) as the key/value pairs (since each text
    # atom could contain info about its source)
    def __init__(self, input_atom, num_overlap_words, delimiter, split_fcn):
        self.num_overlap_words = num_overlap_words
        # TODO: pass in more arguments here, if necessary
        # Break the input TextAtom into a list of words or word fragments,
        # depending on the split_fcn that was supplied by the caller
        self.atom_list = split_fcn(input_atom)
        self.delimiter = delimiter
        
    def __iter__(self):
        self.curr_pos = 0
        return self
        
    def determine_end_pos(self, begin_pos):
        return min(self.num_overlap_words + begin_pos, len(self.atom_list))
    
    def next(self):
        
        # Note that here, begin_pos, end_pos, etc. keep track of which *word*
        # we're currently on
        
        if self.curr_pos == -1:
            raise StopIteration
        
        # If there are fewer remaining words than num_overlap_words, then there's no
        # point in breaking it up since nothing can "attach" to it
        begin_pos = self.curr_pos
        
        # Get the next batch of words - end at num_overlap_words or the end of
        # the source text, whichever comes first
        end_pos = self.determine_end_pos(begin_pos)
        
        # The key is a list of atoms
        key_list = self.atom_list[begin_pos : end_pos]
        
        # advance to the next group of atoms (i.e., the value)     
        begin_pos = self.curr_pos + self.num_overlap_words
        end_pos = self.determine_end_pos(begin_pos)
        
        if end_pos > len(self.atom_list) or end_pos < begin_pos:
            # There aren't enough words remaining to fill a atom, so arrange it
            # so that the next call will stop the iteration, and then return
            # the remaining text as the key and None as the value so that the
            # caller will know to end its random text generation after this
            # atom.
            value_list = None
            self.curr_pos = -1
        else:
            value_list = self.atom_list[begin_pos : end_pos]
            self.curr_pos += 1
        
        # TODO: What we really need to do is concatenate the atoms in key_list
        # into a single atom (same goes for value_list)
        
        # return a tuple containing two TextAtoms, each with num_overlap_words words
        # this means we have to first convert from lists of words to strings
        
        key = atom_list_to_atom(key_list, self.delimiter)
        value = atom_list_to_atom(value_list, self.delimiter)
        return (key, value)
        

class MarkovManager(object):
    """
    Manages a dictionary in which the key is a word sequence that appeared in the
    input text and the value is a list in which each element is a word sequence
    that followed the corresponding key in the input.
    
    Also stores a list of word sequences that can start a chain of words (e.g.,
    in the case of film titles, the elements in this list would be word
    sequences that were found to begin film titles.)  MarkovManager can start
    a new sequence by choosing an element at random from this list.
    """
    
    def __init__(self):
        self.dict = {}
        self.first_words = []
        
    def add(self, first_words, next_words):
        """
        Associates value with key in the dictionary
        """
        
        # If the dictionary already contains an entry for the first_words,
        # then we want to append our new next_words to the next words that
        # already appear in the list.
        if self.dict.has_key(first_words):
            value = self.dict[first_words]
        else:
            value = []
        
        # add the next_words onto the end of the list for this key
        value.append(next_words)
         
        # store the value back in the dictionary
        self.dict[first_words] = value
        
    def add_first_words(self, these_first_words):
        """
        Adds these_first_words to the list of words that are known to 
        start a word sequence
        """
        self.first_words.append(these_first_words)
        
    def inspect(self):
        for key in self.dict.keys():
            print 'key =', key, 'value = ', self.dict[key]
            
    def generate_text(self, delimiter=' ', max_groups=500):
        
        # TODO: Have this generate a *list* of TextGroups (though this list may
        # only have one element) 
        
        # choose an element at random from the list of word sequences that
        # can begin a sequence
        current_atom = random.choice(self.first_words)
        
        # Create the initial TextGroup
        current_text_group = TextGroup()
        
#        result = TextGroup()
        text_group_list = []
        
        for group_num in range(max_groups):

            if current_atom == None:
                # TODO: we probably don't always want to separate output
                # atoms with newlines
#                result += '\n'
                
                # There's no current atom, so we need to close up the current
                # TextGroup, add it to the list, and create a new one to store
                # whatever's coming next
                
                text_group_list.append(current_text_group)
                current_text_group = TextGroup()
                
                # Find the next TextAtom
                current_atom = random.choice(self.first_words)
            
#            result += current_atom + delimiter
            # Add the atom to the TextGroup
            current_text_group.append(current_atom)
#            text_group_list.append(current_atom)
            
            # grab the next group
            if not self.dict.has_key(current_atom):
                raise Exception("Couldn't find key " + current_atom + ' in dictionary')
                            
            current_atom = random.choice(self.dict[current_atom])  

        # finally...
        return text_group_list
    
    def load_text(self, text_atoms, depth, delimiter, split_fcn):
        # iterate through each title and store the pieces in the dictionary
        for text_atom in text_atoms:
            saw_first_words = False
            # for words, use ' ' as the delimiter; for letters, use ''
            for (key, value) in TextualUnitSlicer(text_atom, depth, delimiter,
                                                  split_fcn):
                self.add(key, value)
                if not saw_first_words and not key == None:
                    self.add_first_words(key)
                    saw_first_words = True    
